# -*- coding: utf-8 -*-

import torch.nn.functional as F


def calculate_acc(logit, labels, ignore_index=-100):
    """
    计算准确率的函数，主要用于忽略特定索引的预测结果
    参数:
    logit (Tensor): 模型的输出结果，形状为(batch_size, seq_length, vocab_size)
    labels (Tensor): 实际标签，形状为(batch_size, seq_length)
    ignore_index (int): 需要忽略的索引，默认值为-100
    返回:
    float: 准确率
    """
    # 对模型输出结果进行预处理，移除最后一个时间步的输出，并调整形状为(-1, vocab_size)
    # 示例:[4,95,13317]->[4,94,13317]->[376,13317]
    logit = logit[:, :-1, :].contiguous().view(-1, logit.size(-1))
    # 对标签进行预处理，移除第一个时间步的标签，并调整形状为(-1)
    # 示例:[4,95]->[4,94]->[376]
    labels = labels[:, 1:].contiguous().view(-1)
    # print(f'labels真实标签---》{logit.shape}')
    # 对每个预测单词，取出最大概率值以及对应索引
    _, logit = logit.max(dim=-1)  # 对于每条数据，返回最大的index
    # print(f'logit预测结果---》{logit.shape}')
    '''
    在 PyTorch 中，labels.ne(ignore_index) 表示将标签张量 labels 中的值不等于 ignore_index(-100)
    的位置标记为 True，
    等于 ignore_index(-100) 的位置标记为 False。这个操作，以过滤掉 ignore_index 对损失的贡献
    '''
    non_pad_mask = labels.ne(ignore_index)
    # TODO 以下计算正确预测的数量和总词数
    '''
    在 PyTorch 中，logit.eq(labels) 表示将模型的预测输出值 logit 中等于标签张量 labels 的位置标记为 True，
    不等于标签张量 labels 的位置标记为 False。以标记出预测输出值和标签值相等的位置。
    masked_select(non_pad_mask) 表示将张量中非填充标记的位置选出来。
    '''
    # masked_select(non_pad_mask) 使用non_pad_mask对上述布尔张量进行掩码操作，仅保留非填充位置的结果
    # sum().item() 对保留的结果求和并转换为Python标量，得到正确预测的数量
    n_correct = logit.eq(labels).masked_select(non_pad_mask).sum().item()
    # non_pad_mask.sum().item() 对non_pad_mask求和并转换为Python标量，得到总词数（非填充位置的数量）
    n_word = non_pad_mask.sum().item()
    # 返回正确预测的数量和总词数
    return n_correct, n_word

